#include <cv_bridge/cv_bridge.h>
#include <image_transport/image_transport.h>
#include <message_filters/cache.h>
#include <message_filters/subscriber.h>
#include <ros/ros.h>
#include <std_msgs/Int8.h>

#include <chrono>
#include <mutex>
#include <opencv2/opencv.hpp>
#include <thread>

#include "INIReader.h"
#include "custom_msg/GPS.h"
#include "flashai/commons/spdlog_utils.h"
#include "flashai/commons/string_utils.h"
#include "flashai/commons/visualize.h"
#include "flashros/commons/utils.h"
#include "flashros_commons/BoxesMsg.h"
#include "vision_ruler/VisionRuler.h"

void BoxesMsgConvertion(const flashros_commons::BoxesMsgConstPtr& msg,
                        std::vector<flashai::commons::Box>& boxes) {
  if (msg == nullptr) {
    SPDLOG_DEBUG("NULL boxes msg");
    return;
  }

  for (int i = 0; i < msg->num_objects; i++) {
    flashai::commons::Box box;
    box.x1 = msg->bboxes[i * 4];
    box.y1 = msg->bboxes[i * 4 + 1];
    box.x2 = msg->bboxes[i * 4 + 2];
    box.y2 = msg->bboxes[i * 4 + 3];
    box.score = msg->socres[i];
    box.distance = msg->distances[i];
    box.class_id = msg->classes[i];
    box.text = msg->texts[i];
    boxes.push_back(box);
  }
}

class MsgHelper {
 public:
  MsgHelper() {}

  void Init(std::string mode, ros::NodeHandle& nh, INIReader& reader,
            std::string section) {
    int buffer_size = reader.GetInteger(section, "msg_buffer_size", 1);
    mode_ = mode;

    std::vector<std::string> topic_names;
    flashai::utils::SplitStr(reader.GetString(section, "msg_topics", ""),
                             topic_names);
    SPDLOG_DEBUG("Start Init {} MsgHelper {} msg topics", mode_,
                 topic_names.size());

    for (auto topic_name : topic_names) {
      SPDLOG_DEBUG("Init msg topic {}", topic_name);
      auto sub = new message_filters::Subscriber<flashros_commons::BoxesMsg>(
          nh, topic_name, buffer_size);
      auto cache = new message_filters::Cache<flashros_commons::BoxesMsg>(
          *sub, buffer_size);
      subs_.push_back(sub);
      caches_.push_back(cache);
    }
    SPDLOG_DEBUG("Init msg topic successfully...");

    ruler_.initFromConfigFile(reader, section);
    SPDLOG_DEBUG(
        "Init vision ruler with {} CameraInfos for {} mode successfully...",
        ruler_.cameraInfos.size(), mode_);

    SPDLOG_INFO("Successfully Init {} MsgHelper {} msg topics", mode_,
                topic_names.size());

    SPDLOG_DEBUG("Start to init camera info module");
    // ros subscriber callback function
    auto camera_info_callback_fn = [&](const std_msgs::Int8::ConstPtr& msg) {
      camera_info_lock_.lock();
      cur_camera_info_id_ = msg->data;
      ROS_WARN_STREAM("Change " << mode_ << " camera info id to "
                                << cur_camera_info_id_);
      camera_info_lock_.unlock();
    };
    std::string camera_info_topic_name =
        reader.GetString(section, "camera_info_topic_name", "");
    camera_info_sub_ = nh.subscribe<std_msgs::Int8>(camera_info_topic_name, 1,
                                                    camera_info_callback_fn);
    SPDLOG_DEBUG("Successfully init camera info module");

    ROS_WARN("Successfully init CameraFuse");
    SPDLOG_INFO("Successfully init CameraFuse");
  }

  std::vector<flashai::commons::Box> GetLatextBoxes(
      vision_ruler::global_info ginfo) {
    std::vector<flashai::commons::Box> boxes;
    SPDLOG_DEBUG("Try to get {} boxes from cache", mode_);
    for (auto cache : caches_) {
      auto msg = cache->getElemAfterTime(start_time_);
      BoxesMsgConvertion(msg, boxes);
    }
    SPDLOG_DEBUG("Suceessfully get {} {} boxes from cache", boxes.size(),
                 mode_);
    for (auto& box : boxes) {
      SPDLOG_DEBUG("box {} {} {} {}", box.x1, box.y1, box.x2, box.y2);
    }

    camera_info_lock_.lock();
    int camera_info_id = cur_camera_info_id_;
    camera_info_lock_.unlock();
    if (camera_info_id >= ruler_.cameraInfos.size()) {
      ROS_WARN_STREAM("Unexpect " << mode_ << " camera info id "
                                  << camera_info_id << ", use 0 instead");
      camera_info_id = 0;
    }
    ruler_.computeDistance(boxes, ginfo, camera_info_id);
    SPDLOG_DEBUG("Suceessfully update vision ruler msg for {}", mode_);

    UpdateBoxesTexts(boxes);
    return boxes;
  }

  void Release() {}

 private:
  std::vector<message_filters::Cache<flashros_commons::BoxesMsg>*> caches_;
  std::vector<message_filters::Subscriber<flashros_commons::BoxesMsg>*> subs_;
  std::string mode_;
  ros::Time start_time_ = ros::Time::now();
  vision_ruler::VisionRuler ruler_;
  vision_ruler::global_info default_ginfo_;

  // Camera Info change
  int cur_camera_info_id_ = 0;
  std::mutex camera_info_lock_;
  ros::Subscriber camera_info_sub_;
  int default_camera_info_id_ = 0;

  void UpdateBoxesTexts(std::vector<flashai::commons::Box>& boxes) {
    // TODO: 根据需求，展示所需的内容
    // 这里要做的就是根据条件，修改 box.text 字段
    // 后续可视化内容就是展示该 box.text 字段内容

    // 目前的类别一共分为：
    // 0) 人：显示 `行人`
    // 1）人头：显示 `人头/距离(m)/纬度,经度`
    // 2）机动车：显示 `机动车`
    // 3）车牌OCR：显示 `OCR内容/距离(m)/纬度,经度`
    // 4）普通OCR：显示 `OCR内容`
    // 5）激光测距：显示 `距离(m)/纬度,经度`
    for (auto& box : boxes) {
      std::stringstream ss;
      switch (box.class_id) {
        case 0:
          ss << "行人";
          break;
        case 1:
          ss << "头/" << flashai::utils::FloatToString(box.distance, 2) << "m/"
             << flashai::utils::FloatToString(box.latitude, 6) << ","
             << flashai::utils::FloatToString(box.longitude, 6);
          break;
        case 2:
          ss << "机动车";
          break;
        case 3:
          ss << box.text << box.distance << "m/" << box.latitude << ","
             << box.longitude;
          break;
        case 4:
          ss << box.text;
          break;
        case 5:
          ss << flashai::utils::FloatToString(box.distance, 2) << "m/"
             << flashai::utils::FloatToString(box.latitude, 6) << ","
             << flashai::utils::FloatToString(box.longitude, 6);
          break;
        default:
          ROS_WARN_STREAM("Get box with unexpected class id " << box.class_id);
          break;
      }
      box.text = ss.str();
    }
  }
};

class CameraSubscribe {
 public:
  CameraSubscribe() {}
  void Init(std::string mode, ros::NodeHandle& nh, INIReader& reader,
            std::string section) {
    SPDLOG_DEBUG("Start init CameraSubscribe");
    mode_ = mode;
    max_wait_secs_ =
        reader.GetInteger(section, "camera_max_wait_secs", 5000000);
    running_ = true;
    std::string topic_name = reader.GetString(section, "camera_topic_name", "");

    // ros subscriber callback function
    auto callback_fn = [&](const sensor_msgs::ImageConstPtr& msg) {
      try {
        lock_.lock();
        auto cur_frame = cv_bridge::toCvShare(msg, "bgr8")->image;
        auto cur_time = msg->header.stamp;
        frames_info_.clear();
        frames_info_.push_back(
            std::pair<ros::Time, cv::Mat>(cur_time, cur_frame));
        ROS_INFO_STREAM_THROTTLE(
            25, "Add one frame with shape ("
                    << frames_info_.back().second.rows << ", "
                    << frames_info_.back().second.cols << ") to " << mode_);
        lock_.unlock();
      } catch (cv_bridge::Exception& e) {
        SPDLOG_ERROR("Could not convert from '{}' to 'bgr8'.",
                     msg->encoding.c_str());
      }
    };

    // start a new thread to run subscriber
    SPDLOG_DEBUG("Start init {} cam thread", mode_);
    thread_ = std::thread([&]() {
      image_transport::ImageTransport it(nh);
      image_transport::Subscriber sub =
          it.subscribe(topic_name, 1, callback_fn);
      SPDLOG_INFO("Ros Input thread started with topic {}", topic_name);
      while (running_) {
        ros::spinOnce();
      }
    });
    SPDLOG_DEBUG("Successfully init {} cam thread", mode_);

    std::this_thread::sleep_for(std::chrono::seconds(5));
    SPDLOG_INFO(
        "Ros input pipeline init successfully with topic_name name '{}'",
        topic_name);
  }

  cv::Mat GetLatestFrame() {
    SPDLOG_DEBUG("Try to get {} latest frame", mode_);
    cv::Mat frame;
    lock_.lock();
    SPDLOG_DEBUG("Current image size {}", frames_info_.size());
    SPDLOG_DEBUG("{} {}", frames_info_.back().second.rows,
                 frames_info_.back().second.cols);
    if (frames_info_.size() > 0) {
      frame = frames_info_.back().second.clone();
    }
    SPDLOG_DEBUG("{} {}", frames_info_.back().second.rows,
                 frames_info_.back().second.cols);
    lock_.unlock();
    SPDLOG_DEBUG("Successfully get {} latest frame {} {}", mode_, frame.rows,
                 frame.cols);
    return frame;
  }

  ros::Time GetLatestTime() {
    ros::Time time(0, 0);
    lock_.lock();
    if (frames_info_.size() > 0) {
      time.sec = frames_info_.back().first.sec;
      time.nsec = frames_info_.back().first.nsec;
    }
    lock_.unlock();
    SPDLOG_DEBUG("Successfully get latest time {} {} for {} mode", time.sec,
                 time.nsec, mode_);
    return time;
  }

 private:
  std::string mode_;
  std::vector<std::pair<ros::Time, cv::Mat>> frames_info_;
  std::thread thread_;
  std::mutex lock_;
  bool running_;
  int max_wait_secs_;
};

class CameraFuse {
 public:
  CameraFuse(INIReader& reader, std::string rgb_section, std::string ir_section,
             std::string visualize_section, std::string output_section,
             std::string gps_section, 
             std::string data_root, ros::NodeHandle& nh) {
    // rgb msg and camera
    SPDLOG_DEBUG("Start to init rgb msg and cam");
    rgb_msg_.Init("rgb", nh, reader, rgb_section);
    rgb_cam_.Init("rgb", nh, reader, rgb_section);
    SPDLOG_DEBUG("Successfully init rgb msg and cam");

    // ir msg and camera
    SPDLOG_DEBUG("Start to init ir msg and cam");
    ir_msg_.Init("ir", nh, reader, ir_section);
    ir_cam_.Init("ir", nh, reader, ir_section);
    SPDLOG_DEBUG("Successfully init ir msg and cam");

    visualizer_.Init(reader, visualize_section, data_root);
    SPDLOG_DEBUG("Successfully init visualizer");

    it_ = new image_transport::ImageTransport(nh);
    pub_ = it_->advertise(
        reader.GetString(output_section, "topic_name", ""),
        reader.GetInteger(output_section, "topic_queue_size", 1));
    max_fps_ = reader.GetInteger(output_section, "max_fps", 25);
    SPDLOG_DEBUG("Successfully init Publisher");

    SPDLOG_DEBUG("Start to init gps module");
    int buffer_size = reader.GetInteger(gps_section, "gps_buffer_size", 1);
    gps_sub_ = new message_filters::Subscriber<custom_msg::GPS>(
        nh, reader.GetString(gps_section, "gps_topic_name", ""), buffer_size);
    gps_cache_ =
        new message_filters::Cache<custom_msg::GPS>(*gps_sub_, buffer_size);
    auto gps_callback_fn = [](const custom_msg::GPS::ConstPtr& msg) {
      ROS_WARN_THROTTLE(1, "Get GPS Info %.6f %.6f", msg->lat * 1.0,
                        msg->lon * 1.0);
    };
    gps_cache_->registerCallback(gps_callback_fn);
    camera_gps_degree_ = reader.GetReal(gps_section, "camera_gps_degree", 0);
    laser_gps_degree_ = reader.GetReal(gps_section, "laser_gps_degree", 0);
    SPDLOG_DEBUG("Successfully init gps module");
  }

  void LoopOnce() {
    SPDLOG_DEBUG("Start loop once");
    ros::Time rgb_time = rgb_cam_.GetLatestTime();
    ros::Time ir_time = ir_cam_.GetLatestTime();

    if (rgb_time.sec == 0 && rgb_time.nsec == 0 && ir_time.sec == 0 &&
        ir_time.nsec == 0) {
      SPDLOG_WARN("No input for the loop");
      return;
    }

    // Get Latext GPS info and camera info id
    vision_ruler::global_info ginfo = GetLatextGpsInfo();

    if (rgb_time > ir_time) {
      // rgb mode
      SPDLOG_DEBUG("RGB Mode");

      // draw image
      auto boxes = rgb_msg_.GetLatextBoxes(ginfo);
      SPDLOG_DEBUG("Succesfully get {} boxes", boxes.size());
      for (auto box : boxes) {
        SPDLOG_DEBUG("box: {} {} {}", box.distance, box.latitude,
                     box.longitude);
      }
      auto frame = rgb_cam_.GetLatestFrame();
      visualizer_.Draw(frame, boxes);
      SPDLOG_DEBUG("Succesfully draw image for rgb Mode {} {}", frame.rows,
                   frame.cols);

      // generate msg
      auto msg =
          cv_bridge::CvImage(std_msgs::Header(), "bgr8", frame).toImageMsg();
      msg->header.stamp = ros::Time::now();
      pub_.publish(msg);
      SPDLOG_INFO("Succesfully publish rgb draw image with {} boxes",
                  boxes.size());
    } else {
      // ir mode
      SPDLOG_DEBUG("IR Mode");

      // draw image
      auto boxes = ir_msg_.GetLatextBoxes(ginfo);
      SPDLOG_DEBUG("Succesfully get {} boxes", boxes.size());
      auto frame = ir_cam_.GetLatestFrame();
      visualizer_.Draw(frame, boxes);
      SPDLOG_DEBUG("Succesfully draw image for ir Mode {} {}", frame.rows,
                   frame.cols);

      // generate msg
      auto msg =
          cv_bridge::CvImage(std_msgs::Header(), "bgr8", frame).toImageMsg();
      msg->header.stamp = ros::Time::now();
      pub_.publish(msg);
      SPDLOG_INFO("Succesfully publish ir draw image with {} boxes",
                  boxes.size());
    }
  }

  void Run() {
    SPDLOG_DEBUG("Start run");
    ros::Rate loop_rate(max_fps_);
    int i = 0;
    while (ros::ok()) {
      LoopOnce();
      ros::spinOnce();
      loop_rate.sleep();
    }
  }

  void Release() {
    delete gps_cache_;
    delete gps_sub_;
    delete it_;
    pub_.shutdown();
    visualizer_.Release();
  }

 private:
  // Get all model predictions
  MsgHelper rgb_msg_, ir_msg_;

  // Get RGB and IR raw images
  CameraSubscribe rgb_cam_, ir_cam_;
  image_transport::ImageTransport* it_;
  image_transport::Publisher pub_;

  // gps
  message_filters::Cache<custom_msg::GPS>* gps_cache_;
  ros::Time start_time_ = ros::Time::now();
  message_filters::Subscriber<custom_msg::GPS>* gps_sub_;
  double camera_gps_degree_, laser_gps_degree_;

  flashai::utils::DetectionVisualizer visualizer_;
  int max_fps_;

  vision_ruler::global_info GetLatextGpsInfo() {
    vision_ruler::global_info ginfo;
    auto gps_msg = gps_cache_->getElemAfterTime(start_time_);
    if (gps_msg == nullptr) {
      SPDLOG_WARN("No gps data received...");
    } else {
      if (std::abs(gps_msg->lat) > 1e-5 && std::abs(gps_msg->lon) > 1e-5) {
        ginfo.latitude = gps_msg->lat;
        ginfo.longitude = gps_msg->lon;
        ginfo.altitude = 0;  // 目前没有海拔高度，传入0
        ginfo.pose_pitch = 0;
        ginfo.camera_gps_degree = camera_gps_degree_;
        ginfo.lidar_gps_degree = laser_gps_degree_;
        ginfo.pose_yaw = (gps_msg->gpsAngle) / 180 * PI;
        ginfo.withGps = true;
        SPDLOG_INFO("Get gps data...");
      } else {
        SPDLOG_INFO("Get empty data...");
      }
    }

    return ginfo;
  }
};

int main(int argc, char** argv) {
  ros::init(argc, argv, "fuse_by_cache");
  ros::NodeHandle nh;

  // param
  std::string data_root = flashros::commons::GetRosParam(nh, "data_root");
  ROS_WARN_STREAM("Data root " << data_root);

  // config
  std::string base_config_name = "/config_filename";
  std::string config_file_path = flashros::commons::GetRosParam(
      nh, ros::this_node::getName() + base_config_name);
  flashai::utils::CheckFileExists(config_file_path);
  INIReader cfg_reader(config_file_path);
  ROS_WARN_STREAM("Config file path " << config_file_path);

  // init spdlog
  flashai::utils::InitSpdLogger(cfg_reader);
  SPDLOG_INFO("Successfully init spdlog");

  // prepare for fuse
  std::string rgb_section = "rgb";
  std::string ir_section = "ir";
  std::string visualize_section = "visualize";
  std::string output_section = "output";
  std::string gps_section = "gps";

  CameraFuse fuse(cfg_reader, rgb_section, ir_section, visualize_section,
                  output_section, gps_section, data_root,
                  nh);
  fuse.Run();

  return 0;
}
